#include <iostream>
#include <Eigen/Dense>  // eigen library
#include <ctime>        // clock, etc
#include <iomanip>      // std::setprecision, etc
#include <cmath>        // use fabs, pow, etc
#include <climits>


/* 迭代法解线性方程组
* Ax = b
* 该迭代法收敛的必要条件有两条
*** 条件1：对角线元素不等于零
*** 条件2：0 < w < 2
* 迭代法解线性方程组的充要条件有一个
*** 谱半径 < 1 
* 
*/
enum IterativeSolverMethod{
    JACOBI,               // 雅可比迭代法
    GAUSS_SEIDEL,         // 高斯迭代法
    SOR                   // 超松弛迭代法
};
/* JacobiIterative: 雅可比迭代法解线性方程组
*/
void JacobiIterative(const Eigen::MatrixXd& A,     // 系数矩阵
                     const Eigen::VectorXd& b,      // 常数项 
                     Eigen::VectorXd& x){            // 解
    // 行遍历，每次迭代计算一个解
    Eigen::VectorXd x_temp = Eigen::VectorXd(x.size());
    for(int i = 0; i < A.rows(); ++ i){
        x_temp[i] = (b[i] - (A.row(i).dot(x) - x[i]*A(i, i)))/A(i, i);
    }
    x = x_temp;   // 全部算完再更新
}

/* GaussSeidelIterative: 高斯-赛达尔迭代法，雅克比迭代法的改进版本
* 区别在于，每次迭代计算一个新的解，均采用当前最新的结果
*/
void GaussSeidelIterative(const Eigen::MatrixXd& A,   // 系数矩阵
                          const Eigen::VectorXd& b,        // 常数项 
                          Eigen::VectorXd& x){             // 解
    // 行遍历，每次迭代计算一个解
    for(int i = 0; i < A.rows(); ++ i){
        x[i] = (b[i] - (A.row(i).dot(x) - x[i]*A(i, i)))/A(i, i);   // 实时更新
    }
}

/* SorIterative: 给定超松弛系数的超松弛迭代法
* 当w == 1时，就是Gauss-Seidel迭代法
*/
void SorIterative(const Eigen::MatrixXd& A,      // 系数矩阵
                  const Eigen::VectorXd& b,      // 常数项 
                  const double& w,               // 松弛因子w, 通常满足  0 < w < 2
                  Eigen::VectorXd& x){           // 解
    // 行遍历，每次迭代计算一个解
    for(int i = 0; i < A.rows(); ++ i){
        double x_gauss_seidel = (b[i] - (A.row(i).dot(x) - x[i]*A(i, i)))/A(i, i);  // Gauss-Seidel 迭代的结果
        x[i] = x[i] + w * (x_gauss_seidel - x[i]);
    }    
}
 
/* IterativeSolver: 选择迭代方法，根据阈值迭代求解线性方程组的解
*/
void IterativeSolver(const Eigen::MatrixXd& A,   // 系数矩阵
                     const Eigen::VectorXd& b,      // 常数项 
                     const int& threshold_iteration_times,            // 迭代次数
                     const double& threshold__iteration_error,        // 迭代误差
                     const IterativeSolverMethod& solver,                   // 迭代方法
                     Eigen::VectorXd& x){            // 解
    // x = Eigen::VectorXd::Zero(A.rows(), 1);  // 初始解
    x = Eigen::VectorXd::Zero(A.rows());  // 初始解
    Eigen::VectorXd last_x = x;   // 保存上一次的解用于计算两次迭代之间的误差

    // 判断是否满足迭代条件
    int current_iteration_times = 1;
    double current_iteration_error = static_cast<double>(INT_MAX);      // 借用INT_MAX对初始误差进行初始化

    const double w = 1.3;   // 在switch case外定义
    while(current_iteration_times < threshold_iteration_times &&  current_iteration_error > threshold__iteration_error){
        switch(solver){
            case JACOBI:
                JacobiIterative(A, b, x);
                break;
            case GAUSS_SEIDEL:
                GaussSeidelIterative(A, b, x);
                break;
            case SOR:
                SorIterative(A, b, w, x);
                break;
            default:
                break;
        }
        
        current_iteration_error = (x-last_x).norm();
        last_x = x;
        current_iteration_times ++;
    }
    std::cout << "current_iteration_error = (solution-last_solution).norm() = " << current_iteration_error << std::endl;
    std::cout << "current_iteration_times = " << current_iteration_times << std::endl;
}

void TestIterativeSolver(){
    const int dim = 4;
    Eigen::MatrixXd A(dim, dim);
    Eigen::VectorXd b(dim, 1);
    A << -4, 1, 1, 1,
        1, -4, 1, 1,
        1, 1, -4, 1,
        1, 1, 1, -4;
    b << 1, 1, 1, 1;

    // ldl分解求解精确值
    Eigen::VectorXd x_qr(dim, 1);
    x_qr = A.colPivHouseholderQr().solve(b);

    // 迭代法求解
    const int threshold_iteration_times = 100;        // 迭代次数
    const double threshold__iteration_error = std::pow(0.1, 4);    // 迭代误差

    Eigen::VectorXd x_jacobi_iteration;
    IterativeSolverMethod solver = JACOBI;
    IterativeSolver(A, b, threshold_iteration_times, threshold__iteration_error, solver, x_jacobi_iteration);

    Eigen::VectorXd x_gauss_seidel_iteration;
    solver = GAUSS_SEIDEL;
    IterativeSolver(A, b, threshold_iteration_times, threshold__iteration_error, solver, x_gauss_seidel_iteration);

    Eigen::VectorXd x_sor_iteration;
    solver = SOR;
    IterativeSolver(A, b, threshold_iteration_times, threshold__iteration_error, solver, x_sor_iteration); 

    std::cout << "x_qr = " << x_qr << std::endl;
    std::cout << "x_jacobi_iteration = " << x_jacobi_iteration << std::endl;
    std::cout << "absolute error of jacobi iteration is: " << (x_jacobi_iteration - x_qr).norm() << std::endl; 

    std::cout << "x_gauss_seidel_iteration = " << x_gauss_seidel_iteration << std::endl;
    std::cout << "absolute error of gauss seidel iteration is: " << (x_gauss_seidel_iteration - x_qr).norm() << std::endl; 

    std::cout << "x_sor_iteration = " << x_sor_iteration << std::endl;
    std::cout << "absolute error of SOR iteration is: " << (x_sor_iteration - x_qr).norm() << std::endl;    
}

int main()
{
    TestIterativeSolver();
    system("pause");
    return 0;
}